import os
import shutil
import pandas as pd
from tqdm import tqdm

def split_dataset(base_path='./trainset', output_path='./processed_dataset'):
    """
    Split dataset into real and fake folders based on label.txt
    
    Args:
        base_path: Path to original dataset
        output_path: Path where real and fake folders will be created
    """
    # Read the label file
    label_file = os.path.join(base_path, 'train_label.txt')
    df = pd.read_csv(label_file)
    
    # Create output directories
    real_dir = os.path.join(output_path, 'real')
    fake_dir = os.path.join(output_path, 'fake')
    
    os.makedirs(real_dir, exist_ok=True)
    os.makedirs(fake_dir, exist_ok=True)
    
    # Source directory for images
    img_dir = os.path.join(base_path, 'train')
    
    print("Splitting dataset into real and fake folders...")
    for _, row in tqdm(df.iterrows(), total=len(df)):
        img_name = row['img_name']
        target = row['target']
        
        # Source and destination paths
        src_path = os.path.join(img_dir, img_name)
        dst_path = os.path.join(fake_dir if target == 1 else real_dir, img_name)
        
        # Copy file if it exists
        if os.path.exists(src_path):
            shutil.copy2(src_path, dst_path)
        else:
            print(f"Warning: Image {img_name} not found in {img_dir}")
    
    # Count files in each directory
    real_count = len(os.listdir(real_dir))
    fake_count = len(os.listdir(fake_dir))
    
    print(f"\nDataset split complete:")
    print(f"Real images: {real_count}")
    print(f"Fake images: {fake_count}")
    print(f"Total images: {real_count + fake_count}")
    
    # Create pickle files for compatibility
    real_files = [os.path.join(real_dir, f) for f in os.listdir(real_dir)]
    fake_files = [os.path.join(fake_dir, f) for f in os.listdir(fake_dir)]
    
    import pickle
    
    with open(os.path.join(output_path, 'train.pickle'), 'wb') as f:
        pickle.dump(real_files + fake_files, f)
    
    # Optional: create a validation split
    val_ratio = 0.1
    val_size_per_class = int(min(len(real_files), len(fake_files)) * val_ratio)
    
    val_real = real_files[-val_size_per_class:]
    val_fake = fake_files[-val_size_per_class:]
    
    with open(os.path.join(output_path, 'val.pickle'), 'wb') as f:
        pickle.dump(val_real + val_fake, f)
    
    print(f"\nCreated pickle files:")
    print(f"train.pickle: {len(real_files + fake_files)} images")
    print(f"val.pickle: {len(val_real + val_fake)} images")

if __name__ == '__main__':
    split_dataset()